import os
import numpy as np
import json
import open3d as o3d
from scipy.spatial.transform import Rotation as R


def GetRotationByBoneIndex(bone_rotations, index):
    quat = np.array([bone_rotations[index]["x"], bone_rotations[index]["y"], bone_rotations[index]["z"], bone_rotations[index]["w"]])
    return R.from_quat(quat).as_matrix()


def GetTransformationByBoneIndex(boneinfos, bone_positions, bone_rotations, root_index, index):
    transform = np.eye(4)
    pos = boneinfos[index]["InitPosition"] if index != root_index else bone_positions[index]
    transform[0, 3] = pos["x"]
    transform[1, 3] = pos["y"]
    transform[2, 3] = pos["z"]
    transform[:3, :3] = GetRotationByBoneIndex(bone_rotations, index)
    return transform

        
def compute_global_transform(boneinfos, bone, parent_transform, bone_positions, bone_rotations):
    local_transform = GetTransformationByBoneIndex(boneinfos, bone_positions, bone_rotations, root_index=0, index=bone["Index"])
    global_transform = np.dot(parent_transform, local_transform)
    return global_transform


# 定义一个函数来遍历关节树并计算关节位置
def traverse_skeleton(boneinfos, bone, parent_transform, bone_positions, bone_rotations, joint_poses):
    global_transform = compute_global_transform(boneinfos, bone, parent_transform, bone_positions, bone_rotations)
    joint_poses[bone["Index"]] = global_transform.copy()
    for child in bone["Childs"]:
        traverse_skeleton(boneinfos, boneinfos[child], global_transform, bone_positions, bone_rotations, joint_poses)


def load_roam_date_one_sequence(boneinfos, meta, start_frame=0, end_frame=-1, sampling_rate=1):
    """
    return: a np.float32 representing the pose of each joint in each frame in world space, shape = (N_frame, 27, 4, 4)
    """

    total_positions = json.load(open(meta["human_positions_path"], "r"))
    total_rotations = json.load(open(meta["human_rotations_path"], "r"))

    if end_frame == -1:
        end_frame = len(total_positions) - 1
    assert start_frame <= end_frame

    seq = []
    for frame_idx in range(start_frame, end_frame+1, sampling_rate):

        bone_positions = total_positions[frame_idx]
        bone_rotations = total_rotations[frame_idx]

        final_pose = [np.eye(4) for _ in range(len(boneinfos))]
        traverse_skeleton(boneinfos, boneinfos[0], np.eye(4), bone_positions, bone_rotations, final_pose)
        final_pose = np.float32(final_pose)
        assert final_pose.shape == (27, 4, 4)
        seq.append(final_pose)
    seq = np.float32(seq)

    # eliminate the world space difference between ROAM and H1 environment
    T = np.eye(4).astype(np.float32)
    T[:3, :3] = np.float32([[1,0,0],[0,0,-1],[0,1,0]])
    seq = T.reshape(1, 1, 4, 4) @ seq

    return seq


def load_roam_data_overall(roam_data_metadata, start_frame=0, end_frame=-1, sampling_rate=1):

    boneinfos = json.load(open(roam_data_metadata["human_skeleton_path"], "r"))  # bone info

    roam_sequence_info = []
    for meta in roam_data_metadata["sequences"]:
        seq = load_roam_date_one_sequence(boneinfos, meta, start_frame=start_frame, end_frame=end_frame, sampling_rate=sampling_rate)
        roam_sequence_info.append(seq)

    return roam_sequence_info
